package main

import (
	"fmt"
)

// 定义颜色
const (
	RED   = true
	BLACK = false
)

// 节点结构体
type Node struct {
	value  int
	color  bool
	left   *Node
	right  *Node
	parent *Node
}

// 红黑树结构体
type RedBlackTree struct {
	root *Node
}

// 创建新节点
func NewNode(value int) *Node {
	return &Node{
		value: value,
		color: RED,
	}
}

// 左旋
func (tree *RedBlackTree) LeftRotate(root *Node) {
	newRoot := root.right
	root.right = newRoot.left
	if newRoot.left != nil {
		newRoot.left.parent = root
	}
	newRoot.parent = root.parent
	if root.parent == nil {
		tree.root = newRoot
	} else if root == root.parent.left {
		root.parent.left = newRoot
	} else {
		root.parent.right = newRoot
	}
	newRoot.left = root
	root.parent = newRoot
}

// 右旋
func (tree *RedBlackTree) RightRotate(root *Node) {
	newRoot := root.left
	root.left = newRoot.right
	if newRoot.right != nil {
		newRoot.right.parent = root
	}
	newRoot.parent = root.parent
	if root.parent == nil {
		tree.root = newRoot
	} else if root == root.parent.left {
		root.parent.left = newRoot
	} else {
		root.parent.right = newRoot
	}
	newRoot.right = root
	root.parent = newRoot
}

// 插入修复
func (tree *RedBlackTree) fixInsert(node *Node) {
	for node.parent != nil && node.parent.color == RED {
		if node.parent == node.parent.parent.left {
			uncle := node.parent.parent.right
			if uncle != nil && uncle.color == RED {
				// Case 1: 叔叔节点是红色
				// 将父节点和叔叔节点设为黑色
				node.parent.color = BLACK
				uncle.color = BLACK
				// 将祖父节点设为红色
				node.parent.parent.color = RED
				// 将当前节点指向祖父节点，继续检查
				node = node.parent.parent
			} else {
				if node == node.parent.right {
					// Case 2: 当前节点是父节点的右子节点
					// 先对父节点进行左旋转
					node = node.parent
					tree.LeftRotate(node)
				}
				// 将父节点设为黑色，祖父节点设为红色
				node.parent.color = BLACK
				node.parent.parent.color = RED
				// 对祖父节点进行右旋转
				tree.RightRotate(node.parent.parent)
			}
		} else {
			uncle := node.parent.parent.left
			if uncle != nil && uncle.color == RED {
				// mirror case 1: 叔叔节点是红色
				// 将父节点和叔叔节点设为黑色
				node.parent.color = BLACK
				uncle.color = BLACK
				// 将祖父节点设为红色
				node.parent.parent.color = RED
				// 将当前节点指向祖父节点，继续检查
				node = node.parent.parent
			} else {
				if node == node.parent.left {
					// mirror case 2: 当前节点是父节点的左子节点
					// 先对父节点进行右旋转
					node = node.parent
					tree.RightRotate(node)
				}
				// 将父节点设为黑色，祖父节点设为红色
				node.parent.color = BLACK
				node.parent.parent.color = RED
				// 对祖父节点进行左旋转
				tree.LeftRotate(node.parent.parent)
			}
		}
	}
	// 确保根节点始终为黑色
	tree.root.color = BLACK
}

// 插入节点
func (tree *RedBlackTree) Insert(value int) {
	newNode := NewNode(value)
	var parent *Node = nil
	current := tree.root

	for current != nil {
		parent = current
		if newNode.value < current.value {
			current = current.left
		} else {
			current = current.right
		}
	}
	newNode.parent = parent

	if parent == nil {
		tree.root = newNode
	} else if newNode.value < parent.value {
		parent.left = newNode
	} else {
		parent.right = newNode
	}

	tree.fixInsert(newNode)
}

// 中序遍历
func (tree *RedBlackTree) InOrderTraversal(node *Node) {
	if node != nil {
		tree.InOrderTraversal(node.left)
		fmt.Printf("%d ", node.value)
		tree.InOrderTraversal(node.right)
	}
}

func main() {
	tree := &RedBlackTree{}

	elements := []int{20, 15, 25, 10, 5, 30, 35}
	for _, value := range elements {
		tree.Insert(value)
	}

	fmt.Println("红黑树的中序遍历：")
	tree.InOrderTraversal(tree.root) // 输出: 5 10 15 20 25 30 35
	fmt.Println()
}
